"""
learning rate module
"""
import numpy as np


class _LRScheduler():
    """学习率调度器基类
    定义学习率调度器的基本接口和属性
    """
    def __init__(self, learning_rate, max_epoch, steps_per_epoch):
        self.base_lr = learning_rate       # 基础学习率
        self.steps_per_epoch = steps_per_epoch  # 每个epoch的步数
        self.total_steps = int(max_epoch * steps_per_epoch)  # 总训练步数

    def get_lr(self):
        """计算学习率的抽象方法
        子类需实现该方法以定义具体的学习率变化策略
        """
        raise NotImplementedError


class _WarmUp():
    """学习率预热基类
    定义学习率预热阶段的基本接口
    """
    def __init__(self, warmup_init_lr):
        self.warmup_init_lr = warmup_init_lr  # 预热初始学习率

    def get_lr(self, current_step):
        """获取预热阶段学习率的抽象方法
        Args:
            current_step: 当前训练步数
        """
        raise NotImplementedError


class _LinearWarmUp(_WarmUp):
    """线性预热策略
    学习率从warmup_init_lr线性增加到base_lr
    """
    def __init__(self, lr, warmup_epochs, steps_per_epoch, warmup_init_lr=0):
        self.base_lr = lr                          # 目标学习率
        self.warmup_init_lr = warmup_init_lr       # 预热起始学习率
        self.warmup_steps = int(warmup_epochs * steps_per_epoch)  # 预热总步数
        super(_LinearWarmUp, self).__init__(warmup_init_lr)

    def get_warmup_steps(self):
        """获取预热阶段的总步数"""
        return self.warmup_steps

    def get_lr(self, current_step):
        """计算当前预热步骤的学习率
        Args:
            current_step: 当前预热步骤（从1开始）
        Returns:
            float: 当前学习率
        """
        # 计算线性增长步长
        lr_inc = (float(self.base_lr) - float(self.warmup_init_lr)) / float(self.warmup_steps)
        # 线性增长公式
        learning_rate = float(self.warmup_init_lr) + lr_inc * current_step
        return learning_rate


class StepLR(_LRScheduler):
    """阶梯式学习率调度器
    每经过epoch_size个epoch，学习率乘以gamma衰减

    Args:
        lr (float): 初始学习率
        epoch_size (int): 学习率衰减的周期（epoch数）
        gamma (float): 学习率衰减因子（如0.1表示衰减为1/10）
        steps_per_epoch (int): 每个epoch的训练步数
        max_epoch (int): 总训练epoch数
        warmup_epochs (int): 预热阶段的epoch数，默认为0
    """
    def __init__(self, lr, epoch_size, gamma, steps_per_epoch, max_epoch, warmup_epochs=0):
        self.epoch_size = epoch_size      # 衰减周期（epoch）
        self.gamma = gamma                # 衰减因子
        # 初始化线性预热策略
        self.warmup = _LinearWarmUp(lr, warmup_epochs, steps_per_epoch)
        super(StepLR, self).__init__(lr, max_epoch, steps_per_epoch)

    def get_lr(self):
        """生成整个训练过程的学习率序列
        Returns:
            np.ndarray: 学习率数组，形状为[total_steps]
        """
        warmup_steps = self.warmup.get_warmup_steps()  # 预热阶段步数
        lr_each_step = []

        for i in range(self.total_steps):
            if i < warmup_steps:
                # 预热阶段使用线性增长学习率
                learning_rate = self.warmup.get_lr(i+1)
            else:
                # 计算当前处于第几个epoch
                cur_ep = i // self.steps_per_epoch
                # 阶梯式衰减公式：base_lr * gamma^(cur_ep // epoch_size)
                learning_rate = self.base_lr * self.gamma**(cur_ep // self.epoch_size)

            lr_each_step.append(learning_rate)
        
        return np.array(lr_each_step).astype(np.float32)